import torch.nn as nn
import torch

from ggnn.src.model.layers import Propogator, AttrProxy


class GGNN(nn.Module):
    """
    Gated Graph Sequence Neural Networks (GGNN)
    Mode: SelectNode
    Implementation based on https://arxiv.org/abs/1511.05493
    """

    def __init__(self, opt):
        super(GGNN, self).__init__()
        # 要求h_v维度大于node annotations
        assert (opt.state_dim >= opt.annotation_dim, 'state_dim must be no less than annotation_dim')
        # h_v维度
        self.state_dim = opt.state_dim
        # node annotation维度
        self.annotation_dim = opt.annotation_dim
        # 最大边id
        self.n_edge_types = opt.n_edge_types
        # 最大点id
        self.n_node = opt.n_node
        # GNN层数(parser)
        self.n_steps = opt.n_steps

        # 每种边的类型 + 方向(in/out)独立
        for i in range(self.n_edge_types):
            # incoming and outgoing edge embedding
            # 论文中邻接矩阵A中的每一个cell: B, B', C, C'
            in_fc = nn.Linear(self.state_dim, self.state_dim)
            out_fc = nn.Linear(self.state_dim, self.state_dim)
            self.add_module("in_{}".format(i), in_fc)
            self.add_module("out_{}".format(i), out_fc)

        self.in_fcs = AttrProxy(self, "in_")
        self.out_fcs = AttrProxy(self, "out_")

        # GGNN模型: Propogation Model + Output Model
        # Propogation Model
        # GRU消息传递模型
        self.propogator = Propogator(self.state_dim, self.n_node, self.n_edge_types)

        # Output Model
        # 输出模型, 将h_v与x_v（未padding的）做拼接输入神经网络
        self.out = nn.Sequential(
            # [h_v^T, x_v]
            nn.Linear(self.state_dim + self.annotation_dim, self.state_dim),
            nn.Tanh(),
            nn.Linear(self.state_dim, 1)
        )

        self._initialization()

    def _initialization(self):
        """
        参数初始化
        :return:
        """
        for m in self.modules():
            if isinstance(m, nn.Linear):
                m.weight.data.normal_(0.0, 0.02)
                m.bias.data.fill_(0)

    def forward(self, prop_state, annotation, A):
        """
        前向传播函数
        :param prop_state: 按batch训练，维度N * D
        :param annotation: node annotation x_v(未padding的)
        :param A:
        :return:
        """
        for i_step in range(self.n_steps):
            in_states = []
            out_states = []
            for i in range(self.n_edge_types):
                in_states.append(self.in_fcs[i](prop_state))
                out_states.append(self.out_fcs[i](prop_state))
            # https://zhuanlan.zhihu.com/p/64551412 (解释contiguous)
            # 做拼接
            in_states = torch.stack(in_states).transpose(0, 1).contiguous()
            # -1是bacth的数量, self.n_node*self.n_edge_types表示N * 边的类型数量, state_dim为论文中的h_v的embedding维度为D
            in_states = in_states.view(-1, self.n_node * self.n_edge_types, self.state_dim)
            # 做拼接
            out_states = torch.stack(out_states).transpose(0, 1).contiguous()
            # 同in_states
            out_states = out_states.view(-1, self.n_node * self.n_edge_types, self.state_dim)
            # 用GRU的方式消息传递, 这里理解算邻居消息(论文中的a_v)时,
            # 实际上和RGCN非常相似: a_v = sum_{forall j \in N(i)} W(e)h_j, 就是邻居embedding的相加
            # 这里面的W(e)神经网络参数是和“边的类型 + 边的方向"有关的
            prop_state = self.propogator(in_states, out_states, prop_state, A)

        # 将h_v^T与node annotation x_v(未padding的)拼接起来
        join_state = torch.cat((prop_state, annotation), 2)

        output = self.out(join_state)
        output = output.sum(2)

        return output
